! (C) Copyright 2005- ECMWF.
! 
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
! In applying this licence, ECMWF does not waive the privileges and immunities
! granted to it by virtue of its status as an intergovernmental organisation
! nor does it submit to any jurisdiction.
!

PROGRAM TEST_ADJOINT
USE PARKIND1  ,ONLY : JPIM     ,JPRB
USE MPL_MODULE  ,ONLY : MPL_INIT, MPL_END, MPL_BARRIER, MPL_MYRANK
USE ABORT_TRANS_MOD  ! This is not really correct usage
!

IMPLICIT NONE


INTEGER(KIND=JPIM) :: NSMAX,NDGL,NPROC,NPRGPNS,NPRGPEW,NPRTRW,NPRTRV,NREP
INTEGER(KIND=JPIM) :: NOUT,MYPROC,NSPECG,NSPEC2G,NGPTOTG
INTEGER(KIND=JPIM) :: NFLEV,NFLEVG
INTEGER(KIND=JPIM) :: NSPEC2,NGPTOT,NPROMA,NGPBLKS,MYSETV,NUMP
INTEGER(KIND=JPIM) :: IVSET(1000)
INTEGER(KIND=JPIM) ,ALLOCATABLE :: NLOEN(:),ITO(:),MYMS(:),NASM0(:),IPRCIDS(:)
INTEGER(KIND=JPIM) :: JLEV,JREP,JROC

CHARACTER*6 CLNAME

REAL(KIND=JPRB) , ALLOCATABLE :: ZSPECX(:,:),ZSPECY(:,:),ZSPECP(:,:)
REAL(KIND=JPRB) , ALLOCATABLE :: ZVORX(:,:),ZVORY(:,:),ZVORP(:,:)
REAL(KIND=JPRB) , ALLOCATABLE :: ZDIVX(:,:),ZDIVY(:,:),ZDIVP(:,:)
REAL(KIND=JPRB) , ALLOCATABLE :: ZGX(:,:,:)
REAL(KIND=JPRB) , ALLOCATABLE :: ZSPECYG(:,:),ZSPECXG(:,:)
REAL(KIND=JPRB) , ALLOCATABLE :: ZNORM(:),GW(:),ZRANDSP(:),ZRANDGP(:,:),ZPERT(:,:,:)
REAL(KIND=JPRB) :: ZSC
!REAL(KIND=JPRB) :: ZWALL,ZWALL0
REAL(KIND=JPRB),EXTERNAL :: TIMEF
LOGICAL :: LSTATS,LSTATSCPU,LSYNCSTATS,LDETAILED_STATS,&
            &LSTATS_OMP, LSTATS_COMMS,&
            &LTRACE_STATS, LBARRIER_STATS, LBARRIER_STATS2,&
            &LSTATS_ALLOC, LSTATS_MEM, LXML_STATS, LMPOFF
INTEGER(KIND=JPIM) ::NTRACE_STATS, NPRNT_STATS, NSTATS_MEM
REAL(KIND=JPRB) :: ZAVEAVE(0:200)

#include "setup_trans0.h"
#include "setup_trans.h"
#include "trans_inq.h"
#include "specnorm.h"
#include "dir_trans.h"
#include "inv_trans.h"
#include "dir_transad.h"
#include "inv_transad.h"
#include "dist_grid.h"
#include "gath_spec.h"
#include "dist_spec.h"
#include "gath_grid.h"
#include "gstats_setup.intfb.h"

! ======== PARAMETERS WHICH MAY BE MODIFIED, ORIGINALLY COMING FROM NAMELIST ===========
NDGL = 32
NFLEVG = 9
NPROMA = 8
NPRGPNS = 1
NPRGPEW = 1   ! NPROC must match NPRGPNS * NPRGPEW
NPRTRW = 1
NPRTRV = 1
NREP = 1
! ======================================================================================

LSTATS = .TRUE.
LSTATSCPU = .FALSE.
LSYNCSTATS = .TRUE.
LDETAILED_STATS = .FALSE.
LBARRIER_STATS = .FALSE.
LBARRIER_STATS2 = .FALSE.
LSTATS_ALLOC = .FALSE.
LSTATS_MEM = .FALSE.
LTRACE_STATS = .FALSE.
LXML_STATS = .FALSE.
NTRACE_STATS = 0
NPRNT_STATS = 3
NSTATS_MEM = 0
NPROC = NPRGPNS*NPRGPEW
IF(NPROC /= NPRTRW*NPRTRV) THEN
  PRINT *,'NPRGPNS,NPRGPEW,NPRTRW,NPRTRV ',NPRGPNS,NPRGPEW,NPRTRW,NPRTRV
  CALL ABORT_TRANS('NPRGPNS*NPRGPEW /= NPRTRW*NPRTRV')
ENDIF
NSMAX = (2*NDGL-1)/3
NSPECG = (NSMAX+1)*(NSMAX+2)/2
NSPEC2G = NSPECG*2
IF(NPROC > 1 ) THEN
  CALL MPL_INIT
  MYPROC = MPL_MYRANK()
  NOUT = 20
  WRITE(CLNAME,'(A,I2.2)') 'OUT.',MYPROC
  OPEN(NOUT,FILE=CLNAME)
  LMPOFF = .FALSE.
ELSE
  LMPOFF = .TRUE.
  NOUT = 6
  MYPROC = 1
ENDIF
ALLOCATE(IPRCIDS(NPROC))
DO JROC=1,NPROC
  IPRCIDS(JROC) = JROC
ENDDO
WRITE(NOUT,*) ' NSMAX= ',NSMAX
WRITE(NOUT,*) ' NDGL=',NDGL
WRITE(NOUT,*) ' LMPOFF= ',LMPOFF

MYSETV = MOD(MYPROC-1,NPRTRV)+1
ALLOCATE(ZSPECYG(NFLEVG,NSPEC2G))
ALLOCATE(ZSPECXG(NFLEVG,NSPEC2G))
ALLOCATE(ZRANDSP(NSPEC2G))
ALLOCATE(NLOEN(NDGL))
ALLOCATE(ZNORM(NFLEVG))
ALLOCATE(ITO(NFLEVG))
NLOEN(:) = 2*NDGL
CALL SETUP_TRANS0(KOUT=NOUT,KERR=0,KPRINTLEV=0,KMAX_RESOL=1,&
 & KPRGPNS=NPRGPNS,KPRGPEW=NPRGPEW,KPRTRW=NPRTRW,LDMPOFF=LMPOFF)
CALL SETUP_TRANS(KSMAX=NSMAX,KDGL=NDGL,KLOEN=NLOEN,LDSPLIT=.TRUE.)
CALL TRANS_INQ(KSPEC2=NSPEC2,KGPTOT=NGPTOT,KGPTOTG=NGPTOTG,KNUMP=NUMP)

ALLOCATE(MYMS(NUMP))
ALLOCATE(NASM0(0:NSMAX))
ALLOCATE(GW(NDGL))
ALLOCATE(ZRANDGP(NGPTOTG,1))
CALL TRANS_INQ(KMYMS=MYMS,KASM0=NASM0,PGW=GW)
NGPBLKS = (NGPTOT-1)/NPROMA+1

WRITE(NOUT,*) ' NSPEC2=',NSPEC2
WRITE(NOUT,*) ' NGPTOT=',NGPTOT

NFLEV = 0
DO JLEV=1,NFLEVG
  IVSET(JLEV) = MOD(JLEV,NPRTRV)+1
  IF(IVSET(JLEV) == MYSETV) THEN
    NFLEV = NFLEV+1
  ENDIF
ENDDO
WRITE(NOUT,*)' NFLEV=',NFLEV
WRITE(NOUT,*) 'SETUP FINISHED'
FLUSH(NOUT)

ALLOCATE(ZSPECX(NFLEV,NSPEC2))
ALLOCATE(ZSPECY(NFLEV,NSPEC2))
ALLOCATE(ZSPECP(NFLEV,NSPEC2))
ALLOCATE(ZVORX(NFLEV,NSPEC2))
ALLOCATE(ZVORY(NFLEV,NSPEC2))
ALLOCATE(ZVORP(NFLEV,NSPEC2))
ALLOCATE(ZDIVX(NFLEV,NSPEC2))
ALLOCATE(ZDIVY(NFLEV,NSPEC2))
ALLOCATE(ZDIVP(NFLEV,NSPEC2))
ALLOCATE(ZGX(NPROMA,3*NFLEVG,NGPBLKS))
ALLOCATE(ZPERT(NPROMA,1,NGPBLKS))

! Prepare perturbations

IF(MYPROC == 1) THEN
  DO JLEV=1,NFLEVG
    CALL RANDOM_NUMBER(ZRANDSP)
    ZSPECYG(JLEV,:) = (1.0_JPRB-2.0_JPRB*ZRANDSP(:))
    CALL RANDOM_NUMBER(ZRANDSP)
    ZSPECXG(JLEV,:) = (1.0_JPRB-2.0_JPRB*ZRANDSP(:))
  ENDDO
ENDIF
CALL RANDOM_NUMBER(ZRANDGP)
ZRANDGP(:,:) = (1.0_JPRB-2.0_JPRB*ZRANDGP(:,:))
ITO(:) = 1
ZVORX = 0.0_JPRB
ZVORY = 0.0_JPRB
ZVORP = 0.0_JPRB
ZDIVX = 0.0_JPRB
ZDIVY = 0.0_JPRB
ZDIVP = 0.0_JPRB

! Distribute perturbations

CALL DIST_GRID(PGPG=ZRANDGP,KFDISTG=1,KFROM=ITO,PGP=ZPERT,KPROMA=NPROMA)
CALL DIST_SPEC(PSPECG=ZSPECXG,KFDISTG=NFLEVG,KFROM=ITO,PSPEC=ZSPECX,&
 & KVSET=IVSET(1:NFLEVG))

!ZWALL=TIMEF()
!ZWALL0=ZWALL
CALL GSTATS(0,0)
CALL GSTATS_SETUP(NPROC,MYPROC,IPRCIDS,&
 &LSTATS,LSTATSCPU,LSYNCSTATS,LDETAILED_STATS,LBARRIER_STATS,LBARRIER_STATS2,&
 &LSTATS_OMP,LSTATS_COMMS,LSTATS_MEM,NSTATS_MEM,LSTATS_ALLOC,&
 &LTRACE_STATS,NTRACE_STATS,NPRNT_STATS,LXML_STATS)

CALL GSTATS_PSUT
DO JREP=1,NREP

  CALL DIST_SPEC(PSPECG=ZSPECYG,KFDISTG=NFLEVG,KFROM=ITO,PSPEC=ZSPECY,&
   & KVSET=IVSET(1:NFLEVG))
! Direct caclulations

  CALL INV_TRANS(PSPSCALAR=ZSPECX,PSPVOR=ZVORX,PSPDIV=ZDIVX,PGP=ZGX,&
   & KPROMA=NPROMA,KVSETSC=IVSET(1:NFLEVG), KVSETUV=IVSET(1:NFLEVG))
  CALL GPC(ZGX)
  CALL DIR_TRANS(PSPSCALAR=ZSPECP,PSPVOR=ZVORP,PSPDIV=ZDIVP,PGP=ZGX,&
   &KPROMA=NPROMA,KVSETSC=IVSET(1:NFLEVG), KVSETUV=IVSET(1:NFLEVG))
  CALL SCALPRODSP(ZSPECP,ZSPECY,ZSC)
  IF(JREP == NREP) WRITE(NOUT,*)'<Fx,y> ',ZSC

! Adjoint

  ZSPECP = 0.0_JPRB
  CALL DIR_TRANSAD(PSPSCALAR=ZSPECY,PSPVOR=ZVORY,PSPDIV=ZDIVY,PGP=ZGX,&
   &KPROMA=NPROMA, KVSETSC=IVSET(1:NFLEVG), KVSETUV=IVSET(1:NFLEVG))
  CALL GPCAD(ZGX)
  CALL INV_TRANSAD(PSPSCALAR=ZSPECP,PSPVOR=ZVORP,PSPDIV=ZDIVP,PGP=ZGX,&
   & KPROMA=NPROMA,KVSETSC=IVSET(1:NFLEVG),KVSETUV=IVSET(1:NFLEVG))

  CALL SCALPRODSP(ZSPECX,ZSPECP,ZSC)
  IF(JREP == NREP) WRITE(NOUT,*)'<x,F*y> ',ZSC
!  ZWALL=TIMEF()
!  IF(MYPROC == 1) WRITE(0,*) 'Time spent in loop(i)=',jrep,(zwall-zwall0)*0.001
!  ZWALL0 = ZWALL
ENDDO
CALL GSTATS(0,1)

!ZWALL=TIMEF()
!WRITE(NOUT,*) 'Time spent in main loop=',zwall*0.001
!IF(MYPROC == 1) WRITE(0,*) 'Time spent in main loop=',zwall*0.001

CALL GSTATS_PRINT(NOUT,ZAVEAVE,200)
IF(NPROC > 1 ) THEN
 CALL MPL_BARRIER()
 CALL MPL_END
ENDIF

CONTAINS

SUBROUTINE GPC(PGP)

! Grid-point computations

REAL(KIND=JPRB) :: PGP(:,:,:)

INTEGER(KIND=JPIM) :: JLEV,JKGLO,JROF,IEND,IBL,IOFF

CALL GSTATS(108,0)
!$OMP PARALLEL DO SCHEDULE(STATIC,1) PRIVATE(JLEV,JKGLO,IEND,IOFF,IBL,JROF)
DO JLEV=1,NFLEVG
  DO JKGLO=1,NGPTOT,NPROMA
    IEND = MIN(NPROMA,NGPTOT-JKGLO+1)
    IOFF = JKGLO-1
    IBL  = (JKGLO-1)/NPROMA+1
    DO JROF=1,IEND
      PGP(JROF,JLEV,IBL) = PGP(JROF,JLEV,IBL)*ZPERT(JROF,1,IBL)
    ENDDO
  ENDDO
ENDDO
!$OMP END PARALLEL DO
CALL GSTATS(108,1)


END SUBROUTINE GPC

SUBROUTINE GPCAD(PGP)

! Adjoint of GPC

REAL(KIND=JPRB) :: PGP(:,:,:)
INTEGER(KIND=JPIM) :: JLEV,JKGLO,JROF,IEND,IBL,IOFF

CALL GSTATS(135,0)
!$OMP PARALLEL DO SCHEDULE(STATIC,1) PRIVATE(JLEV,JKGLO,IEND,IOFF,IBL,JROF)
DO JLEV=1,NFLEVG
  DO JKGLO=1,NGPTOT,NPROMA
    IEND = MIN(NPROMA,NGPTOT-JKGLO+1)
    IOFF = JKGLO-1
    IBL  = (JKGLO-1)/NPROMA+1
    DO JROF=1,IEND
      PGP(JROF,JLEV,IBL) = PGP(JROF,JLEV,IBL)*ZPERT(JROF,1,IBL)
    ENDDO
  ENDDO
ENDDO
!$OMP END PARALLEL DO
CALL GSTATS(135,1)


END SUBROUTINE GPCAD

SUBROUTINE SCALPRODSP(PSP1,PSP2,PSC)

! Scalar product in spectral space
REAL(KIND=JPRB) :: PSP1(:,:),PSP2(:,:)
REAL(KIND=JPRB) :: PSC

INTEGER(KIND=JPIM) :: JMLOC,IM,JIR,JN,INM,JLEV
REAL(KIND=JPRB) :: ZMFACT,ZSP(NFLEV,NSPEC2),ZSPG(NFLEVG,NSPEC2G)

PSC = 0.0_JPRB
ZSP(:,:) = 0.0_JPRB

CALL GSTATS(110,0)
!$OMP PARALLEL DO SCHEDULE(STATIC,1) PRIVATE(JLEV,JMLOC,IM,ZMFACT,JIR,JN,INM)
DO JLEV=1,NFLEV
  DO JMLOC=1,NUMP
    IM = MYMS(JMLOC)
    ZMFACT=1.0_JPRB+REAL(MIN(1,IM),JPRB)
    DO JIR=0,MIN(1,IM)
      DO JN=IM,NSMAX
        INM = NASM0(IM)+(JN-IM)*2+JIR
        ZSP(JLEV,INM) = PSP1(JLEV,INM)*PSP2(JLEV,INM)*ZMFACT/2.0_JPRB
      ENDDO
    ENDDO
  ENDDO
ENDDO
!$OMP END PARALLEL DO

CALL GATH_SPEC(PSPECG=ZSPG,KFGATHG=NFLEVG,KTO=ITO,PSPEC=ZSP,&
 & KVSET=IVSET(1:NFLEVG))

IF(MYPROC == 1) THEN
  PSC = SUM(ZSPG)
ELSE
  PSC = 0.0_JPRB
ENDIF
CALL GSTATS(110,1)

END SUBROUTINE SCALPRODSP

END PROGRAM TEST_ADJOINT
